﻿using FlyRemote.Base;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;

namespace FlyRemote.Remote
{
    public enum SnapshotRetrievalMode
    {
        All,

        /// <summary>
        /// 基于过滤器
        /// </summary>
        FromPrefilter,

        /// <summary>
        /// 基于设置创建快照
        /// </summary>
        FromSettings,

        /// <summary>
        /// 基于手动设置
        /// </summary>
        FromManualSettings,


        /// <summary>
        /// 基于用户模式地址创建快照
        /// </summary>
        FromUserModeMemory,

        /// <summary>
        /// 查找堆的快照
        /// </summary>
        FromHeaps,

        /// <summary>
        /// 查找栈空间的快照
        /// </summary>
        FromStacks,

        /// <summary>
        /// 从模块地址创建
        /// </summary>
        FromModules,


        /// <summary>
        /// 从主模块地址
        /// </summary>
        FromExe,


        /// <summary>
        /// 主模块 .data段
        /// </summary>
        FromExeData,


        /// <summary>
        /// 模块地址  .data段
        /// </summary>
        FromModulesData,
    }



    public static class SnapshotMaker
    {
        #region 与设置相关

        /// <summary>
        /// 获取 设置中 要添加的内存保护类型
        /// </summary>
        private static MemoryProtectionEnum GetRequiredProtectionSettings(ScanSettin settin)
        {
            MemoryProtectionEnum result = 0;

            if (settin.RequiredWrite)
            {
                result |= MemoryProtectionEnum.Write;
            }

            if (settin.RequiredExecute)
            {
                result |= MemoryProtectionEnum.Execute;
            }

            if (settin.RequiredCopyOnWrite)
            {
                result |= MemoryProtectionEnum.CopyOnWrite;
            }

            return result;
        }

        /// <summary>
        /// 获取 设置中 被排除的 内存保护类型
        /// </summary>
        private static MemoryProtectionEnum GetExcludedProtectionSettings(ScanSettin settin)
        {
            MemoryProtectionEnum result = 0;

            if (settin.ExcludedWrite)
            {
                result |= MemoryProtectionEnum.Write;
            }

            if (settin.ExcludedExecute)
            {
                result |= MemoryProtectionEnum.Execute;
            }

            if (settin.ExcludedCopyOnWrite)
            {
                result |= MemoryProtectionEnum.CopyOnWrite;
            }

            return result;
        }

        /// <summary>
        /// 取设置  的内存类型
        /// </summary>
        private static MemoryTypeEnum GetAllowedTypeSettings()
        {
            MemoryTypeEnum result = 0;

            if (ScanSettings.Default.MemoryTypeNone)
            {
                result |= MemoryTypeEnum.None;
            }

            if (ScanSettings.Default.MemoryTypePrivate)
            {
                result |= MemoryTypeEnum.Private;
            }

            if (ScanSettings.Default.MemoryTypeImage)
            {
                result |= MemoryTypeEnum.Image;
            }

            if (ScanSettings.Default.MemoryTypeMapped)
            {
                result |= MemoryTypeEnum.Mapped;
            }

            return result;
        }





        #endregion

        #region 各种快照的创建


        private static Snapshot FromVirtualPages(DataType type, MemoryProtectionEnum requiredPageFlags, MemoryProtectionEnum excludedPageFlags, MemoryTypeEnum allowedTypeFlags, UInt64 startAddress, UInt64 endAddress, int alignment = 0)
        {


            List<ReadGroup> memoryRegions = new List<ReadGroup>(0x1000);

            RMFactory.Default.EnumVirtualPages(
                requiredPageFlags,
                excludedPageFlags,
                allowedTypeFlags,
                startAddress,
                endAddress, new VirtualPageEnumProc((ulong address, uint size) =>
                {
                    memoryRegions.Add(ReadGroupManager.Create(type, address, (int)size, alignment));
                    return false;
                }));

            return new Snapshot(null, memoryRegions);
        }


        /// <summary>
        /// 基于设置创建快照
        /// </summary>
        internal static Snapshot FromSettings(DataType type, ScanSettin settin)
        {



            MemoryProtectionEnum requiredPageFlags = GetRequiredProtectionSettings(settin);
            MemoryProtectionEnum excludedPageFlags = GetExcludedProtectionSettings(settin);
            MemoryTypeEnum allowedTypeFlags = GetAllowedTypeSettings();

            UInt64 startAddress;
            UInt64 endAddress;

            if (ScanSettings.Default.IsUserMode)
            {
                startAddress = 0;
                endAddress = RMFactory.Default.GetMaxUsermodeAddress();
            }
            else
            {
                startAddress = ScanSettings.Default.StartAddress;
                endAddress = ScanSettings.Default.EndAddress;
            }

            return FromVirtualPages(type, requiredPageFlags, excludedPageFlags, allowedTypeFlags, startAddress, endAddress, settin.Alignment);

        }


        /// <summary>
        /// 取内存纪录中用到了哪些内存块
        /// </summary>
        static Snapshot GetFilterSnapshot(HashSet<ulong> regions)
        {

            // 整理一个 内存结构的 快照 ,  用来取 内存结构体中的 指针  
            Snapshot heaps = FromHeaps(DataType.Empty);
            HashSet<ReadGroup> readGroups = new HashSet<ReadGroup>();

            foreach (var r in regions)
            {
                bool found = false;

                foreach (var u in heaps.ReadGroups)
                {
                    if (u.BaseAddress <= r && r < u.EndAddress)
                    {
                        readGroups.Add(u);
                        found = true;
                        // 可以直接退出, 处理下一条记录
                        break;
                    }

                }

                if (!found)
                {
                    // 确实有一块内存 状态为保留的 ,  无法匹配到 
                    ;
                    // throw new Exception("堆内存未能命中");
                }

            }

            // 从68块 内存区域 , 精简到 10块 
            return new Snapshot("filer", readGroups.ToList());


        }




        /// <summary>
        /// 创建一个用户空间的快照
        /// </summary>
        static Snapshot FromUsermodeMemory(DataType type)
        {
            MemoryProtectionEnum requiredPageFlags = 0;
            MemoryProtectionEnum excludedPageFlags = 0;
            MemoryTypeEnum allowedTypeFlags = MemoryTypeEnum.None | MemoryTypeEnum.Private | MemoryTypeEnum.Image;

            UInt64 startAddress = 0;
            UInt64 endAddress = RMFactory.Default.GetMaxUsermodeAddress();

            return FromVirtualPages(type, requiredPageFlags, excludedPageFlags, allowedTypeFlags, startAddress, endAddress);

        }


        internal static Snapshot FromAll(DataType type)
        {
            MemoryProtectionEnum requiredPageFlags = 0;
            MemoryProtectionEnum excludedPageFlags = 0;
            MemoryTypeEnum allowedTypeFlags = MemoryTypeEnum.None | MemoryTypeEnum.Private | MemoryTypeEnum.Image | MemoryTypeEnum.Mapped;

            return FromVirtualPages(type, requiredPageFlags, excludedPageFlags, allowedTypeFlags, 0, ulong.MaxValue);

        }

        /// <summary>
        /// 获取 线程栈的快照
        /// </summary>
        internal static Snapshot FromStacks(DataType type)
        {
            List<ReadGroup> memoryRegions = new List<ReadGroup>();
            IEnumerable<NormalizedRegion> virtualPages = RMFactory.Default.GetStackAddresses();

            foreach (NormalizedRegion virtualPage in virtualPages)
            {
                memoryRegions.Add(ReadGroupManager.Create(type, virtualPage.BaseAddress, virtualPage.RegionSize, ScanSettings.Default.Alignment));
            }
            return new Snapshot(null, memoryRegions);
        }

        /// <summary>
        /// 获取所有模块的内存快照
        /// </summary>
        internal static Snapshot FromModules(DataType type)
        {
            IList<ReadGroup> moduleGroups = RMFactory.Default.GetModules().Select(region => ReadGroupManager.Create(type, region.BaseAddress, region.RegionSize, ScanSettings.Default.Alignment, true)).ToList();
            Snapshot moduleSnapshot = new Snapshot(null, moduleGroups);

            return moduleSnapshot;
        }

        /// <summary>
        /// 获取某一个模块的内存快照
        /// </summary>
        internal static Snapshot FromModule(DataType type, string name)
        {
            var list = RMFactory.Default.GetModules();
            IList<ReadGroup> moduleGroups = list.Select(region => region.Name.Contains(name) ? ReadGroupManager.Create(type, region.BaseAddress, region.RegionSize, ScanSettings.Default.Alignment, true) : null).ToList();
            moduleGroups = moduleGroups.Where(u => u != null).ToList();

            Snapshot moduleSnapshot = new Snapshot(null, moduleGroups);
            return moduleSnapshot;
        }


        /// <summary>
        /// 创建一个堆的快照 , 过滤掉模块 和栈 
        /// </summary>
        internal static Snapshot FromHeaps(DataType type)
        {
            IList<NormalizedRegion> expects = RMFactory.Default.GetStackAddresses().ToList();
            IList<NormalizedModule> modules = RMFactory.Default.GetModules().ToList();

            modules.ForEach(o => expects.Add((NormalizedRegion)o));
            expects = expects.OrderBy(o => o.BaseAddress).ToList();

            MemoryProtectionEnum requiredPageFlags = MemoryProtectionEnum.Write;
            MemoryProtectionEnum excludedPageFlags = MemoryProtectionEnum.Execute;
            MemoryTypeEnum allowedTypeFlags = MemoryTypeEnum.None | MemoryTypeEnum.Private;

            UInt64 startAddress = 0;
            UInt64 endAddress = RMFactory.Default.GetMaxUsermodeAddress();

            List<ReadGroup> memoryRegions = new List<ReadGroup>(0x1000);

            int expectIndex = 0;

            RMFactory.Default.EnumVirtualPages(
                requiredPageFlags,
                excludedPageFlags,
                allowedTypeFlags,
                startAddress,
                endAddress, new VirtualPageEnumProc((ulong address, uint size) =>
                {
                    if (expectIndex < expects.Count)
                    {
                        // 符合排除的条件 
                        if (expects[expectIndex].BaseAddress <= address && (address + size) <= expects[expectIndex].BaseAddress)
                        {
                            //如果尾部相等 , 即将 处理下一个 
                            if (address + size == expects[expectIndex].EndAddress)
                            {
                                expectIndex++;
                            }
                            return false;
                        }
                    }

                    memoryRegions.Add(ReadGroupManager.Create(type, address, (int)size, ScanSettings.Default.Alignment));
                    return false;
                }));


            return new Snapshot(null, memoryRegions.ToList());
        }




        /// <summary>
        /// 取可读写内存中 ， 与模块的并集 ， 即为 .Data 段内存 
        /// </summary>
        internal static Snapshot FromModuleData(DataType type, string name = null)
        {
            var list = RMFactory.Default.GetModules();

            if (!string.IsNullOrEmpty(name))
                list = list.Where(region => region.Name.Contains(name));

            return FromModulesData(type, list.ToList());
        }


        static Snapshot FromModulesData(DataType type, IList<NormalizedModule> modules = null)
        {

            if (modules == null)
                modules = RMFactory.Default.GetModules().ToList();

            /// 这里直接 排除系统dll 

            modules = modules.Where(o => !o.FullPath.Contains("C:\\WINDOWS\\")).ToList();

            /// 必须是可写的
            MemoryProtectionEnum requiredPageFlags = MemoryProtectionEnum.Write;
            MemoryProtectionEnum excludedPageFlags = 0;
            MemoryTypeEnum allowedTypeFlags = MemoryTypeEnum.Image;

            UInt64 startAddress = 0;

            ///!!! 对于 64位程序，地址必须够大，否则无法获取到 
            UInt64 endAddress = ulong.MaxValue;

            List<ReadGroup> memoryRegions = new List<ReadGroup>(0x100);


            int expectIndex = 0;

            RMFactory.Default.EnumVirtualPages(
                requiredPageFlags,
                excludedPageFlags,
                allowedTypeFlags,
                startAddress,
                endAddress, new VirtualPageEnumProc((ulong address, uint size) =>
                {
                    if (expectIndex < modules.Count)
                    {
                        // 符合包含的条件 
                        if (modules[expectIndex].BaseAddress < address && modules[expectIndex].BaseAddress <= (address + size))
                        {
                            memoryRegions.Add(ReadGroupManager.Create(type, address, (int)size, ScanSettings.Default.Alignment, true));

                            //如果尾部相等 , 即将 比对下一个 
                            if (address + size == modules[expectIndex].EndAddress)
                            {
                                expectIndex++;
                            }
                            return false;
                        }
                    }

                    return false;
                }));

            return new Snapshot(null, memoryRegions.ToList());
        }


        #endregion


        private static Snapshot _heaps;

        public static Snapshot Create(DataType type, SnapshotRetrievalMode mode)
        {
            switch (mode)
            {
                case SnapshotRetrievalMode.All:
                    return FromAll(type);


                case SnapshotRetrievalMode.FromHeaps:
                    {
                        if (_heaps == null)
                            _heaps = FromHeaps(type);

                        foreach (var u in _heaps.ReadGroups)
                        {
                            u.ElementDataType = type;
                        }
                        return _heaps;
                    }

                case SnapshotRetrievalMode.FromStacks:
                    return FromStacks(type);

                case SnapshotRetrievalMode.FromModules:
                    return FromModules(type);

                case SnapshotRetrievalMode.FromModulesData:
                    return FromModuleData(type);

                case SnapshotRetrievalMode.FromExe:
                    return FromModule(type, PAFactory.Default.OpenedProcess.ProcessName);

                case SnapshotRetrievalMode.FromExeData:
                    return FromModuleData(type, PAFactory.Default.OpenedProcess.ProcessName);

                case SnapshotRetrievalMode.FromManualSettings:
                    return FromSettings(type, ScanSettings.Manual);

                default:
                    return FromSettings(type, ScanSettings.Default);
            }
        }

    }


    public static class SnapshotManager
    {
        /// 这里不合理 , 首先快照数据肯定是 越扫越少的 , 这里如果限制大小
        /// 


        private const int CountLimit = 5;


        static SnapshotManager()
        {
            AccessLock = new Object();
            ObserverLock = new Object();
            Snapshots = new Deque<Snapshot>();
            DeletedSnapshots = new Stack<Snapshot>();
        }

        /// <summary>
        /// 将快照组 改为 双向队列, 允许在 快照过多时, 删除太旧的 快照数据 , 防止战胜过多内存 
        /// </summary>
        public static Deque<Snapshot> Snapshots { get; private set; }

        /// <summary>
        /// Gets the deleted snapshots for the capability of redoing after undo.
        /// </summary>
        public static Stack<Snapshot> DeletedSnapshots { get; private set; }

        /// <summary>
        /// Gets or sets a lock to ensure multiple entities do not try and update the snapshot list at the same time.
        /// </summary>
        private static Object AccessLock { get; set; }

        /// <summary>
        /// Gets or sets a lock to ensure multiple entities do not try and update the snapshot list at the same time.
        /// </summary>
        private static Object ObserverLock { get; set; }



        /// <summary>
        /// 有则返回，没有返回空
        /// </summary>
        public static Snapshot GetActiveSnapshot()
        {
            lock (AccessLock)
            {
                // Take a snapshot if there are none, or the current one is empty
                // 就是因为 搜索到结果是0 , 所以才需要 恢复快照 
                if (Snapshots.Count == 0 || Snapshots.Last() == null)
                {
                    return null;
                }

                // Return the snapshot
                return Snapshots.Last();
            }
        }

        /// <summary>
        /// 恢复上一次撤销的快照
        /// </summary>
        public static void RedoSnapshot()
        {
            lock (AccessLock)
            {
                if (DeletedSnapshots.Count == 0)
                {
                    return;
                }

                Snapshots.AddBack(DeletedSnapshots.Pop());
            }
        }

        /// <summary>
        /// 撤销一个快照
        /// </summary>
        public static void UndoSnapshot()
        {
            lock (AccessLock)
            {
                if (Snapshots.Count == 0)
                {
                    return;
                }

                DeletedSnapshots.Push(Snapshots.RemoveBack());

                if (DeletedSnapshots.Peek() == null)
                {
                    DeletedSnapshots.Pop();
                }

            }
        }

        /// <summary>
        /// Clears all snapshot records.
        /// </summary>
        public static void ClearSnapshots()
        {
            lock (AccessLock)
            {
                Snapshots.Clear();
                DeletedSnapshots.Clear();

                // There can be multiple GB of deleted snapshots, so run the garbage collector ASAP for a performance boost.
                Task.Run(() => GC.Collect());
            }
        }

        /// <summary>
        /// 保存快照， 这将产生一个通知操作
        /// </summary>
        public static void SaveSnapshot(Snapshot snapshot)
        {
            lock (AccessLock)
            {
                // Remove null snapshot if exists
                if (Snapshots.Count != 0 && Snapshots.Last() == null)
                {
                    Snapshots.Last();
                }

                // Do not keep large snapshots in the undo history
                //if (Snapshots.Count != 0 && Snapshots.Last() != null)
                //{
                //    Snapshots.RemoveBack();
                //}

                /// 新的快照在后面, 所以 移除旧的, 从前部删除即可 
                if (Snapshots.Count >= CountLimit)
                    Snapshots.RemoveFront();


                Snapshots.AddBack(snapshot);
                DeletedSnapshots.Clear();
            }
        }




    }

}
